import os
from typing import Optional
from pathlib import Path
from setuptools import setup, Extension
from setuptools.command.build_ext import build_ext
from Cython.Build import cythonize
import numpy as np


class ExtensionOptions:
    def __init__(self,
                 include_numpy: bool = True,
                 avx: bool = False
                 ):
        self.include_numpy: bool = include_numpy
        self.avx: bool = avx


class CustomBuildExt(build_ext):
    def get_ext_filename(self, fullname: str) -> str:
        """ 移除生成pyd文件名中的目标平台内容

        :param fullname: 文件加载名称
        :return: 移除目标平台内容后的pyd路径


        ::

            # 示例1
            result = get_ext_filename("xinfer_yolo.detect.preprocess_cpy.cpy")
            result == "xinfer_yolo/detect/preprocess_cpy.cpy" # true
        """
        original_filename = super().get_ext_filename(fullname).strip()

        if not original_filename.endswith('.pyd'):
            return original_filename

        dirname, filename = os.path.split(original_filename)
        parts = filename.split('.')
        if len(parts) == 3:
            # 移除filename中间的内容
            filename = '{}.{}'.format(parts[0], parts[2])

        filepath = os.path.join(dirname, filename)
        return filepath

    def run(self) -> None:
        """ 删除所有由.pyx生成的.c文件
        """
        project_path = Path(__file__).parent.absolute()
        for ext in self.extensions:
            for src in ext.sources:
                if not src.endswith('.c'):
                    break
                # src为.c文件路径, 例如:"xinfer_yolo/detect/preprocess_cpy.c"
                filepath = project_path / src
                os.remove(str(filepath))
                print(f' * remove: {filepath}')


def create_extension(name: Optional[str] = None,
                     source: Optional[str] = None,
                     option: ExtensionOptions = ExtensionOptions()) -> Extension:
    # todo 完善错误处理
    if source is None:
        raise RuntimeError
    if name is None:
        name = source.replace('/', '.').strip()
        name = name.replace('.pyx', '')

    include_dirs = []
    if option.include_numpy:
        include_dirs.append(np.get_include())

    return Extension(
        name=name,
        sources=[source],
        include_dirs=include_dirs,
        language='3str',
        extra_compile_args=['/arch:AVX'] if option.avx else None
    )


extensions = [
    create_extension(source='xinfer_yolo/detect/preprocess_cpy.pyx'),
    create_extension(source='xinfer_yolo/detect/postprocess_cpy.pyx'),
    create_extension(source='xinfer_yolo/simd/avx_cpy.pyx', option=ExtensionOptions(avx=True))
]

setup(
    name='xinfer_yolo_cpy',
    version='0.0.1',
    ext_modules=cythonize(
        extensions,
    ),
    cmdclass={'build_ext': CustomBuildExt}
)
